package com.demo.userprofile.component.service.democode;

import java.util.List;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.ml.classification.LogisticRegression;
import org.apache.spark.ml.classification.LogisticRegressionModel;
import org.apache.spark.ml.feature.LabeledPoint;
import org.apache.spark.ml.linalg.Vectors;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SQLContext;

import com.google.common.collect.Lists;

/**
 * 挖掘类标签示例代码
 *
 * @author userprofile_demo
 */
public class SparkMlDemoCode {

    public static void main() {
        SparkConf conf = new SparkConf().setAppName("ifMarryOrNotDemo");
        JavaSparkContext jsc = new JavaSparkContext(conf);
        SQLContext jsql = new SQLContext(jsc);

        // 准备训练数据
        List<LabeledPoint> localTraining = Lists.newArrayList(
                new LabeledPoint(1.0, Vectors.dense(0.0, 1.1, 0.1)),
                new LabeledPoint(0.0, Vectors.dense(2.0, 1.0, -1.0)),
                new LabeledPoint(0.0, Vectors.dense(2.0, 1.3, 1.0)),
                new LabeledPoint(1.0, Vectors.dense(0.0, 1.2, -0.5)));
        Dataset<Row> training = jsql.applySchema(jsc.parallelize(localTraining), LabeledPoint.class);

        // 创建LogisticRegression模型示例
        LogisticRegression lr = new LogisticRegression();
        // 设置模型参数
        lr.setMaxIter(10).setRegParam(0.01);
        // 训练模型
        LogisticRegressionModel model1 = lr.fit(training);

        // 准备测试数据
        List<LabeledPoint> localTest = Lists.newArrayList(
                new LabeledPoint(1.0, Vectors.dense(-1.0, 1.5, 1.3)),
                new LabeledPoint(0.0, Vectors.dense(3.0, 2.0, -0.1)),
                new LabeledPoint(1.0, Vectors.dense(0.0, 2.2, -1.5)));
        Dataset<Row> test = jsql.applySchema(jsc.parallelize(localTest), LabeledPoint.class);

        // 测试模型并将结果写入到results中
        model1.transform(test).registerTempTable("results");
        // 打印测试结果
        Dataset<Row> results = jsql.sql("SELECT features, label, probability, prediction FROM results");
        for (Row r : results.collectAsList()) {
            System.out.println("(" + r.get(0) + ", " + r.get(1) + ") -> prob=" + r.get(2) + ", prediction=" + r.get(3));
        }
    }
}
